#! /usr/bin/env python
from sys import argv, exit
import netCDF4
from netCDF4 import Dataset
import numpy as np
import matplotlib
from pylab import *
from matplotlib import colors, ticker, cm
# Sequential colormaps
from matplotlib.cm import Greys, OrRd, YlGnBu, YlOrRd
# Diverging colormaps (white at midpoint)
from matplotlib.cm import Spectral, coolwarm, RdYlBu, RdYlBu_r, RdBu, RdBu_r
# Miscellaneous colormaps
from matplotlib.cm import rainbow, brg, jet, hsv

from matplotlib.colors import LogNorm
from matplotlib import rcParams
#matplotlib.use('agg')
import isFloat
import isInt
execfile('/Users/dkweis/bin/isFloat.py')
execfile('/Users/dkweis/bin/isInt.py')
rcParams.update({'figure.autolayout': True})
print(matplotlib.cm.cmap_d.keys())
#cm=coolwarm()
#cm=winter()

#tit0 = "CESM Model, 5 Tg/yr SO$_2$ Points"
tit0 = "CESM Model, 5 Tg/yr H$_2$SO$_4$ Points"
#tit0 = "CESM Model, Background"

#scenario = "Control"
scenario = "H2SO4_5Tg_points"
#scenario = "H2SO4_5Tg_region"
#scenario = "SO2_5Tg_points"
#scenario = "SO2_5Tg_region"
bgscenario = "Control"

scensuffix = ".001_zm_8yr.nc"
scenprefix = ["so4_a1_CESM2_","so4_a2_CESM2_","so4_a3_CESM2_"]
vble = ["so4_a1","so4_a2","so4_a3"]
# convert from kg(so4)/kg(air) to kg(S)/kg(air)
convert = 3.0
# scale from kg/kg to ppbm
scale = 1.E-9

#print '#args=',len(argv)
#print 'args: ',argv
monthlabels = ['Jan','Feb','Mar','Apr','May','Jun','Jul','Aug','Sep','Oct','Nov','Dec']
latlabels = ['90$^o$S','60$^o$S','30$^o$S','Eq','30$^o$N','60$^o$N','90$^o$N']
lat_vbl = "lat"
lev_vbl = "lev_p"
time_vbl = "day"
lat1 = -90.
lat2 = 90.
alt1 = 0.
alt2 = 35.
altinc = 10.
yticvals = []

nextalt = alt1
while nextalt <= alt2:
    yticvals.append(nextalt)
    nextalt=nextalt + altinc
unit = 0
ndays = 1
days = [0]
ibin = 0
ncontours = 0
contourinc=0.
contourmax=0.
contourmin=0.
avg = 0
tit = ""
iarg=1
while iarg < len(argv):
#    print argv[iarg]
    if argv[iarg] == "-lat":
        iarg += 1
        lat1=float(argv[iarg])
        iarg += 1
        lat2=float(argv[iarg])
        print "found latitude range: ",lat1,lat2
    elif argv[iarg] == "-alt":
        iarg += 1
        alt1=float(argv[iarg])
        iarg += 1
        alt2=float(argv[iarg])
        if (alt2-alt1) < 50.: altinc = 5.
        yticvals = []
        nextalt = alt1
        while nextalt <= alt2:
            yticvals.append(nextalt)
            nextalt=nextalt + altinc
        print "found altitude range: ",alt1,alt2
        print "yticks:  ",yticvals
    elif argv[iarg] == "-vlat":  #latitude variable
        iarg += 1
        lat_vbl = argv[iarg]
        print "found latitude variable: ",lat_vbl
    elif argv[iarg] == "-unit":  
# unit = 0 for mixing ratio plot, unit = 1 for number density plot, unit = 3 for nanobar plot
        iarg += 1
        unit_vbl = argv[iarg]
        if unit_vbl == "molec": unit = 1
        if unit_vbl == "bar": unit = 2
        print "found latitude variable: ",lat_vbl
    elif argv[iarg] == "-vlev":  #variable to plot
        iarg += 1
        lev_vbl = argv[iarg]
        print "found level variable: ",lev_vbl
    elif argv[iarg] == "-vtime":  #variable to plot
        iarg += 1
        time_vbl = argv[iarg]
        print "found time variable: ",time_vbl
    elif argv[iarg] == "-v":  #variable to plot
        iarg += 1
        vble = argv[iarg]
        if vble.find("ODepth") >= 0:
            print "Use colplt for optical depth"
            exit(1)
        print "found variable: ",vble
    elif argv[iarg] == "-title":  #title
        iarg += 1
        tit0 = argv[iarg]
        print "found title: ",tit0
    elif argv[iarg] == "-avg":  #plot annual average
        avg = 1
        print "plotting annual average"
    elif argv[iarg] == "-s":
        iarg += 1
        scale=float(argv[iarg])
        print "found scale factor: ",scale
    elif argv[iarg] == "-cinc":
        iarg += 1
        contourmin = float(argv[iarg])
        iarg += 1
        contourmax = float(argv[iarg])
        iarg += 1
        contourinc = float(argv[iarg])
        print "found contour limits of ",contourmin," - ",contourmax," by ",contourinc
    elif argv[iarg] == "-c":
        contours = []
        iarg += 1
        while iarg < len(argv):
            cntr = argv[iarg]
            if isFloat(cntr):
                contours.append(float(cntr))
                iarg += 1
            else:
                iarg -= 1
                break
        if len(contours) == 3:
            contourmin = contours[0]
            contourmax = contours[1]
            contourinc = contours[2]
            print "found contour limits of ",contourmin," - ",contourmax," by ",contourinc
        else:
            ncontours = len(contours)
            print "found number of coutours = ",ncontours
            print 'contours: ',contours
    elif argv[iarg] == "-bin":
        iarg += 1
        ibin=int(argv[iarg]) - 1
        print "found bin number: ",ibin+1
    elif argv[iarg] == "-n":
        iarg += 1
        ndays=int(argv[iarg])
        print "found number of days: ",ndays
    elif argv[iarg] == "-d":
        days = []
        iarg += 1
        while iarg < len(argv):
            dayz = argv[iarg]
            if isInt(dayz):
                days.append(int(dayz))
                iarg += 1
            else:
                iarg -= 1
                break
        ndays = len(days)
        print "found number of days = ",ndays
        print "found days: ",days
    else:
        scenario = argv[iarg]
        print "found scenario: ",scenario
    iarg += 1



for iv,vbl in enumerate(vble):
    filename = scenprefix[iv] + scenario + scensuffix
    print "opening file ",filename
    ncfile = Dataset(filename,'r')
    bgfilename = scenprefix[iv] + bgscenario + scensuffix
    print "opening file ",bgfilename
    bgfile = Dataset(bgfilename,'r')

    lats = np.array(ncfile.variables[lat_vbl])
    nlats = len(lats)
    if iv==0: print 'lats=',lats[:]
    levs = np.array(ncfile.variables[lev_vbl])
    nlevs = len(levs)
    if iv==0: print 'levs=',levs[:]
    print "Getting data for variable ",vbl
    data1 = np.array(ncfile.variables[vbl])
    print "data1 = ",data1[:,81]
    bg1 = np.array(bgfile.variables[vbl])
    print "bg1 = ",bg1[:,81]
    if iv==0: data = zeros((nlevs,nlats),dtype=float)
    data = data + data1 - bg1
print "Total SO4 = ",data[:,81]
# convert units from kg(SO4) to kg(S)
data = data/convert    
print "Total S = ",data[:,81]
datadims = data.ndim

if unit == 1: 
    temper = np.array(ncfile.variables['TEMPERATURE'])
    airvd = zeros((nlevs,nlats),dtype=float)
    for lat in range(nlats):
        airvd[:,lat]=levs[:]/(temper[:,lat]*1.38E-19)
    data1[:,:] = data1[:,:]*airvd[:,:]
    if nvbles >1: data2[:,:] = data2[:,:]*airvd[:,:]
if unit == 2:
    for lat in range(nlats):
        data1[:,lat] = data1[:,lat]*levs[:]
        if nvbles > 1: data2[:,lat] = data2[:,lat]*levs[:]


x=lats
y=7.*log(1000./levs)
print 'x=',x
print 'y=',y
#if ndays > 990: 
#    ndays = len(time)
#    days = day

if contourinc > 0:
    print "found coutour limits of ",contourmin,contourmax,contourinc
    ncontours = int((contourmax+0.1*contourinc-contourmin)/contourinc) + 1
    contours = []
    for i in range(ncontours):
        cont = contourmin + (i)*contourinc
        contours.append(cont)
    print 'contours: ',contours
if ncontours > 0:
    midlev = int(ncontours/2)
    midcontour = contours[midlev]
    print "Contours: ",contours
    print "contours low, mid, hi: ",contours[0],contours[midlev],contours[-1]

pltrows = 1
pltcols = 1

figure(figsize=(7,6))
axlimits = [lat1,lat2,alt1,alt2]
print "axlimits=",axlimits
xticks(linspace(-90.,90.,7),latlabels,fontsize=14)
yticks(fontsize=14)
xlabel("Latitude",fontsize=16)
ylabel("Approx. Altitude (km)",fontsize=16)
#for month in xrange(0,len(time),step):
nplt = 0
if datadims != 2:
    print "error: program termination"
    exit()
index=-1
count = 0
z = data/scale
axis(axlimits)
tit=tit0
print tit
title(tit,fontsize=14)
minval = 1.e30
maxval = -1.e30
for i in range(nlats):
    for j in range(nlevs):
        if x[i] >= lat1 and x[i] <= lat2:
            if y[j] >= alt1 and y[j] <= alt2:
                minval = min(minval, z[j,i])
                maxval = max(maxval, z[j,i])
print 'min value=',minval,'  max value=',maxval
#cm=winter()
#cm=summer()
if ncontours > 0:
    CS1 = contour(x,y,z,contours)
#        clabel(CS1, colors='black', fmt='%4.2f', fontsize=14)
    clabel(CS1, colors='black', fmt='%4.1f', fontsize=12)
#    CS2 = contourf(x,y,z,contours,colorbar_spacing='proportional')
#    CS2 = contourf(x,y,z,contours,colorbar_spacing='uniform')
    CS2 = contourf(x,y,z,contours,colorbar_spacing='uniform',cmap=get_cmap(OrRd))
#        CS2 = contourf(x,y,z,contours,norm=colors.LogNorm())
    cbar = colorbar(CS2,format='%4.1f')
    cbar.set_ticks(contours)
#        cbar.set_ticklabels(size=14)
else:
    CS1 = contour(x,y,z)
#    clabel(CS1, colors='black', fmt='%4.1f', fontsize=12)
#    CS2 = contourf(x,y,z,colorbar_spacing='proportional')
    CS2 = contourf(x,y,z,colorbar_spacing='uniform', cmap=get_cmap(OrRd))
#        CS2 = contourf(x,y,z,norm=colors.LogNorm())
    colorbar(CS2,format='%4.1f')
#    clabel(CS, inline=1, fontsize=10)

show()
#savefig('contour.pdf')
ncfile.close()
